"""
基于BoT-SORT的多目标跟踪模块
"""

import numpy as np
import cv2
import logging
from typing import List, Dict, Any
from scipy.optimize import linear_sum_assignment
from filterpy.kalman import KalmanFilter

class Track:
    """单个目标轨迹"""
    
    def __init__(self, track_id, initial_detection, config):
        self.track_id = track_id
        self.config = config
        self.hits = 1
        self.age = 0
        self.time_since_update = 0
        
        # 轨迹状态
        self.detection_history = [initial_detection]
        self.features = []  # Re-ID特征
        
        # 初始化卡尔曼滤波器
        self.kf = self.init_kalman_filter(initial_detection)
        
        # 轨迹状态
        self.state = 'tentative'
        
    def init_kalman_filter(self, detection):
        """初始化卡尔曼滤波器"""
        kf = KalmanFilter(dim_x=7, dim_z=4)
        
        # 状态向量: [x, y, w, h, vx, vy, vw]
        # 观测向量: [x, y, w, h]
        
        # 状态转移矩阵
        kf.F = np.array([
            [1, 0, 0, 0, 1, 0, 0],
            [0, 1, 0, 0, 0, 1, 0],
            [0, 0, 1, 0, 0, 0, 1],
            [0, 0, 0, 1, 0, 0, 0],
            [0, 0, 0, 0, 1, 0, 0],
            [0, 0, 0, 0, 0, 1, 0],
            [0, 0, 0, 0, 0, 0, 1]
        ])
        
        # 观测矩阵
        kf.H = np.array([
            [1, 0, 0, 0, 0, 0, 0],
            [0, 1, 0, 0, 0, 0, 0],
            [0, 0, 1, 0, 0, 0, 0],
            [0, 0, 0, 1, 0, 0, 0]
        ])
        
        # 初始化状态
        bbox = detection['bbox']
        kf.x[:4] = [bbox[0], bbox[1], bbox[2], bbox[3]]
        
        # 协方差矩阵
        kf.P[4:, 4:] *= 1000.  # 给速度初值高不确定性
        kf.P *= 10.
        
        # 过程噪声
        kf.Q[4:, 4:] *= 0.01
        
        # 观测噪声
        kf.R[2:, 2:] *= 10.
        
        return kf
    
    def predict(self):
        """预测下一时刻状态"""
        self.kf.predict()
        self.age += 1
        self.time_since_update += 1
        
        # 预测后的边界框
        x, y, w, h = self.kf.x[:4]
        return [x, y, w, h]
    
    def update(self, detection):
        """用检测结果更新轨迹"""
        bbox = detection['bbox']
        measurement = np.array([bbox[0], bbox[1], bbox[2], bbox[3]])
        
        self.kf.update(measurement)
        self.detection_history.append(detection)
        self.hits += 1
        self.time_since_update = 0
        
        # 更新轨迹状态
        if self.hits >= self.config.min_hits:
            self.state = 'confirmed'
        
        # 保存Re-ID特征
        if 'feature' in detection:
            self.features.append(detection['feature'])
            # 只保留最近的特征
            if len(self.features) > self.config.track_buffer:
                self.features.pop(0)
    
    def get_state(self):
        """获取当前状态"""
        x, y, w, h = self.kf.x[:4]
        vx, vy, vw = self.kf.x[4:]
        
        return {
            'track_id': self.track_id,
            'bbox': [x, y, w, h],
            'velocity': [vx, vy],
            'hits': self.hits,
            'age': self.age,
            'state': self.state,
            'last_detection': self.detection_history[-1] if self.detection_history else None,
            'trajectory': [det['bbox'][:2] for det in self.detection_history]
        }

class BotSORT_Tracker:
    """BoT-SORT多目标跟踪器"""
    
    def __init__(self, config):
        self.config = config
        self.logger = logging.getLogger(__name__)
        
        # 跟踪器状态
        self.tracks = []
        self.next_id = 1
        self.frame_count = 0
        
        # Re-ID模型
        self.reid_model = self.load_reid_model()
    
    def load_reid_model(self):
        """加载Re-ID模型"""
        # 这里可以使用预训练的Re-ID模型
        # 例如: torchreid.models.build_model()
        self.logger.info("Re-ID模型加载完成")
        return None
    
    def extract_features(self, detections, frame):
        """提取Re-ID特征"""
        features = []
        for det in detections:
            # 简化版的特征提取
            # 实际应用中应该使用预训练的Re-ID模型
            bbox = det['bbox']
            x, y, w, h = bbox[:4]
            
            # 裁剪车辆区域
            x1, y1 = int(x - w/2), int(y - h/2)
            x2, y2 = int(x + w/2), int(y + h/2)
            
            # 确保坐标在图像范围内
            x1 = max(0, x1); y1 = max(0, y1)
            x2 = min(frame.shape[1], x2); y2 = min(frame.shape[0], y2)
            
            if x2 > x1 and y2 > y1:
                crop = frame[y1:y2, x1:x2]
                # 这里可以添加真正的Re-ID特征提取
                feature = np.random.randn(512)  #  placeholder
                features.append(feature)
            else:
                features.append(np.zeros(512))
                
        return features
    
    def motion_affinity(self, tracks, detections):
        """计算运动亲和度（马氏距离）"""
        cost_matrix = np.zeros((len(detections), len(tracks)))
        
        for i, det in enumerate(detections):
            det_pos = np.array(det['bbox'][:2])
            
            for j, track in enumerate(tracks):
                # 预测位置
                pred_pos = np.array(track.predict()[:2])
                
                # 计算马氏距离
                innovation = det_pos - pred_pos
                # 简化版的马氏距离计算
                mahalanobis_dist = np.sqrt(np.sum(innovation**2))
                
                cost_matrix[i, j] = mahalanobis_dist
                
        return cost_matrix
    
    def appearance_affinity(self, det_features, tracks):
        """计算外观亲和度（余弦相似度）"""
        cost_matrix = np.zeros((len(det_features), len(tracks)))
        
        for i, det_feat in enumerate(det_features):
            for j, track in enumerate(tracks):
                if not track.features:
                    cost_matrix[i, j] = 0.5  # 默认值
                    continue
                    
                # 计算与轨迹特征的余弦相似度
                similarities = []
                for track_feat in track.features:
                    similarity = np.dot(det_feat, track_feat) / (
                        np.linalg.norm(det_feat) * np.linalg.norm(track_feat))
                    similarities.append(similarity)
                
                # 取最大相似度
                max_similarity = max(similarities) if similarities else 0
                cost_matrix[i, j] = 1 - max_similarity  # 转换为代价
                
        return cost_matrix
    
    def iou_affinity(self, detections, tracks):
        """计算IoU亲和度"""
        cost_matrix = np.zeros((len(detections), len(tracks)))
        
        for i, det in enumerate(detections):
            det_corners = det['corners']
            
            for j, track in enumerate(tracks):
                if not track.detection_history:
                    cost_matrix[i, j] = 1.0
                    continue
                    
                # 获取轨迹的最新检测框
                last_det = track.detection_history[-1]
                track_corners = last_det['corners']
                
                # 计算旋转框IoU
                try:
                    rect1 = cv2.minAreaRect(np.array(det_corners, dtype=np.float32))
                    rect2 = cv2.minAreaRect(np.array(track_corners, dtype=np.float32))
                    
                    inter_area, _ = cv2.rotatedRectangleIntersection(rect1, rect2)
                    if inter_area == 0:
                        cost_matrix[i, j] = 1.0
                        continue
                        
                    area1 = cv2.contourArea(cv2.boxPoints(rect1))
                    area2 = cv2.contourArea(cv2.boxPoints(rect2))
                    
                    iou = inter_area / (area1 + area2 - inter_area)
                    cost_matrix[i, j] = 1 - iou
                    
                except:
                    cost_matrix[i, j] = 1.0
                    
        return cost_matrix
    
    def data_association(self, detections, det_features, frame):
        """数据关联"""
        if not self.tracks:
            return [], list(range(len(detections))), []
            
        # 计算各种亲和度
        motion_cost = self.motion_affinity(self.tracks, detections)
        appearance_cost = self.appearance_affinity(det_features, self.tracks)
        iou_cost = self.iou_affinity(detections, self.tracks)
        
        # 综合代价矩阵
        cost_matrix = (
            0.4 * motion_cost +
            0.4 * appearance_cost +
            0.2 * iou_cost
        )
        
        # 匈牙利算法匹配
        det_indices, track_indices = linear_sum_assignment(cost_matrix)
        
        # 筛选有效匹配
        matched_pairs = []
        unmatched_detections = list(range(len(detections)))
        unmatched_tracks = list(range(len(self.tracks)))
        
        for det_idx, track_idx in zip(det_indices, track_indices):
            if cost_matrix[det_idx, track_idx] < self.config.match_threshold:
                matched_pairs.append((det_idx, track_idx))
                if det_idx in unmatched_detections:
                    unmatched_detections.remove(det_idx)
                if track_idx in unmatched_tracks:
                    unmatched_tracks.remove(track_idx)
        
        return matched_pairs, unmatched_detections, unmatched_tracks
    
    def update(self, detections, frame):
        """更新跟踪器"""
        self.frame_count += 1
        
        # 提取Re-ID特征
        det_features = self.extract_features(detections, frame)
        
        # 为检测添加特征
        for i, det in enumerate(detections):
            det['feature'] = det_features[i]
        
        # 预测现有轨迹
        for track in self.tracks:
            track.predict()
        
        # 数据关联
        matched_pairs, unmatched_dets, unmatched_trks = self.data_association(
            detections, det_features, frame)
        
        # 更新匹配的轨迹
        for det_idx, track_idx in matched_pairs:
            self.tracks[track_idx].update(detections[det_idx])
        
        # 创建新轨迹
        for det_idx in unmatched_dets:
            if len(detections[det_idx]['bbox']) >= 4:
                new_track = Track(self.next_id, detections[det_idx], self.config)
                self.tracks.append(new_track)
                self.next_id += 1
        
        # 移除丢失的轨迹
        self.tracks = [track for track in self.tracks 
                      if track.time_since_update <= self.config.max_age]
        
        # 获取跟踪结果
        tracking_results = []
        for track in self.tracks:
            if track.state == 'confirmed':
                tracking_results.append(track.get_state())
        
        return tracking_results